#pragma once
#include <functional>

using namespace std::placeholders;

typedef BOOL(WINAPI* PFN_Wow64DisableWow64FsRedir)(PVOID* OldValue);
typedef BOOL(WINAPI* PFN_Wow64RevertWow64FsRedir)(PVOID OlValue);


class WindowTask
{
public:
	typedef std::function<int(PVOID)>       FuncTask;
	typedef std::function<void(PVOID, int)> FuncRslt;

protected:
	typedef struct _TASK_ITEM
	{
		PVOID       pParam;
		FuncTask    fnTask;
		FuncRslt    fnRslt;
		int         iResult;

		WindowTask* pThis;

	} TASK_ITEM, * PTASK_ITEM;


	//////////////////////////////////////////////////////////////////////////
	// interface
public:
	WindowTask(ULONG uTaskID = WM_USER + 128)
	{
		m_hWindow = NULL;
		m_heExit = NULL;

		m_uErrCtl = 0;
		m_uTskMsg = uTaskID;
		m_uTskCnt = 0;
		m_bAbort = FALSE;
	}


public:
	static int DisableRedir()
	{
		int         iError;
		HMODULE     hKernel;
		FARPROC     pfnProc;
		BOOL        bIsWow64 = FALSE;
		PVOID       pWowData = NULL;

		PFN_Wow64DisableWow64FsRedir    fnWow64DisableWow64FsRedir;
		//PFN_Wow64RevertWow64FsRedir     fnWow64RevertWow64FsRedir;

		do
		{
			iError = IsWow64Process(GetCurrentProcess(), &bIsWow64);
			_ASSERT(iError);
			if (!bIsWow64) { break; }

			hKernel = GetModuleHandle(_T("kernel32.dll"));
			_ASSERT(hKernel);
			pfnProc = GetProcAddress(hKernel, "Wow64DisableWow64FsRedirection");
			_ASSERT(pfnProc);
			if (!pfnProc) { break; }

			fnWow64DisableWow64FsRedir = (PFN_Wow64DisableWow64FsRedir)pfnProc;
			iError = fnWow64DisableWow64FsRedir(&pWowData);
			_ASSERT(iError);

		} while (FALSE);

		return iError;
	}


	int Attach(HWND hwTask, ULONG uErrID)
	{
		m_hWindow = hwTask;
		m_heExit = CreateEvent(NULL, TRUE, TRUE, NULL);

		m_uErrCtl = uErrID;
		m_uTskCnt = 0;

		return 0;
	}


	int PostTask(PVOID pParam, FuncTask fnTask, FuncRslt fnResult)
	{
		int         iError;
		LONG        lCount;
		PTASK_ITEM  ptiTask = NULL;

		_ASSERT(m_hWindow);

		do
		{
			ptiTask = new TASK_ITEM();
			if (!ptiTask) { iError = -ERROR_INSUFFICIENT_BUFFER; break; }

			ptiTask->pParam = pParam;
			ptiTask->fnTask = fnTask;
			ptiTask->fnRslt = fnResult;
			ptiTask->iResult = 0;
			ptiTask->pThis = this;

			ResetEvent(m_heExit);
			InterlockedIncrement((LONG*)& m_uTskCnt);

			iError = QueueUserWorkItem(ThreadTask, ptiTask, 0);
			if (!iError) { iError = -(int)GetLastError(); break; }

			ptiTask = NULL;

		} while (FALSE);

		if (ptiTask)
		{
			lCount = InterlockedDecrement((LONG*)& m_uTskCnt);
			if (0 == lCount) { SetEvent(m_heExit); }
			delete ptiTask;
		}

		return iError;
	}


	int CancelTask(ULONG uTimeout)
	{
		int     iError;

		m_bAbort = TRUE;
		iError = WaitForSingleObject(m_heExit, uTimeout);
		m_bAbort = FALSE;

		return iError;
	}


	//////////////////////////////////////////////////////////////////////////
	// property
protected:
	HWND            m_hWindow;
	ULONG           m_uErrCtl;

	HANDLE          m_heExit;
	ULONG           m_bAbort;
	ULONG           m_uTskMsg;
	ULONG           m_uTskCnt;

	//////////////////////////////////////////////////////////////////////////
	// member
protected:
	BOOL ProcessWindowMessage(HWND hWnd, UINT uMsg, WPARAM wParam,
		LPARAM lParam, LRESULT& lResult)
	{
		BOOL    bHandle = FALSE;

		if (WM_CTLCOLORDLG == uMsg || WM_CTLCOLORSTATIC == uMsg)
		{
			if (m_hWindow && m_uErrCtl &&
				(HWND)lParam == GetDlgItem(m_hWindow, m_uErrCtl))
			{
				HDC     hDC;
				HBRUSH  hBackg = NULL;

				hDC = (HDC)wParam;
				SetTextColor(hDC, RGB(0xFF, 0, 0));
				SetBkMode(hDC, TRANSPARENT);
				hBackg = (HBRUSH)::GetStockObject(NULL_BRUSH);

				lResult = (LRESULT)hBackg;
				bHandle = TRUE;
			}
		}
		else if (m_uTskMsg == uMsg)
		{
			LONG        lTskCnt;
			PTASK_ITEM  ptiTask = (PTASK_ITEM)lParam;

			ptiTask->fnRslt(ptiTask->pParam, ptiTask->iResult);
			delete ptiTask; bHandle = TRUE;

			lTskCnt = InterlockedDecrement((LONG*)& m_uTskCnt);
			if (0 == lTskCnt) { SetEvent(m_heExit); }
		}

		return bHandle;
	}


	static DWORD WINAPI ThreadTask(PVOID pParam)
	{
		int         iError;
		LONG        lTskCnt;
		PTASK_ITEM  ptiTask = (PTASK_ITEM)pParam;
		WindowTask* pThis = (WindowTask*)ptiTask->pThis;

		iError = ptiTask->fnTask(ptiTask->pParam);

		ptiTask->iResult = iError;

		if (pThis->m_bAbort)
		{
			lTskCnt = InterlockedDecrement((LONG*)& pThis->m_uTskCnt);
			if (0 == lTskCnt) { SetEvent(pThis->m_heExit); }
		}
		else
		{
			::SendMessage(pThis->m_hWindow, pThis->m_uTskMsg, 0, (LPARAM)pParam);
		}

		return 0;
	}
};

